# _*_ coding : utf-8 _*_
# @Time : 2025/5/4 16:35
# @Author : liuchengjie
# @File : MQTTDriver
# @Project : esp32
import time
import ubinascii
import json
import machine
import network
from umqtt.simple import MQTTClient


class MQTTDriver:
    def __init__(self, server, port=1883, client_id=None, user=None, password=None,
                 ssl=False, ssl_params=None, keepalive=60):
        """
        初始化MQTT客户端驱动

        参数:
        server -- MQTT服务器地址
        port -- MQTT服务器端口，默认为1883
        client_id -- 客户端ID，不指定则使用设备唯一ID
        user -- MQTT服务器用户名，如果需要认证
        password -- MQTT服务器密码，如果需要认证
        ssl -- 是否使用SSL/TLS连接
        ssl_params -- SSL参数，当ssl=True时使用
        keepalive -- keepalive时间，默认60秒
        """
        # 如果未指定客户端ID，使用设备MAC地址生成唯一ID
        if client_id is None:
            client_id = ubinascii.hexlify(machine.unique_id()).decode()

        self.client_id = client_id
        self.server = server
        self.port = port
        self.user = user
        self.password = password
        self.ssl = ssl
        self.ssl_params = ssl_params
        self.keepalive = keepalive

        self.mqtt_client = None
        self.is_connected = False
        self.last_message = None
        self.subscriptions = {}

    def connect_wifi(self, ssid, password, timeout=10):
        """
        连接到WiFi网络

        参数:
        ssid -- WiFi名称
        password -- WiFi密码
        timeout -- 连接超时时间(秒)

        返回:
        bool -- 是否成功连接
        """
        print(f"正在连接到WiFi网络: {ssid}...")

        wlan = network.WLAN(network.STA_IF)
        wlan.active(True)

        if wlan.isconnected():
            print("已连接到WiFi")
            return True

        wlan.connect(ssid, password)

        # 等待连接或超时
        start_time = time.time()
        while not wlan.isconnected():
            if time.time() - start_time > timeout:
                print("WiFi连接超时！")
                return False
            time.sleep(0.1)

        print("WiFi连接成功")
        print(f"IP地址: {wlan.ifconfig()[0]}")
        return True

    def connect(self):
        """
        连接到MQTT服务器

        返回:
        bool -- 是否成功连接
        """
        try:
            print(f"正在连接到MQTT服务器: {self.server}...")

            # 创建MQTT客户端实例
            self.mqtt_client = MQTTClient(
                self.client_id,
                self.server,
                self.port,
                self.user,
                self.password,
                self.keepalive,
                ssl=self.ssl,
                ssl_params=self.ssl_params
            )

            # 设置回调函数
            self.mqtt_client.set_callback(self._on_message)

            # 连接到服务器
            self.mqtt_client.connect()
            self.is_connected = True

            # 重新订阅之前的主题
            for topic, callback in self.subscriptions.items():
                self.mqtt_client.subscribe(topic)

            print("MQTT服务器连接成功")
            return True

        except Exception as e:
            print(f"MQTT连接失败: {e}")
            self.is_connected = False
            return False

    def disconnect(self):
        """断开MQTT连接"""
        if self.is_connected and self.mqtt_client:
            try:
                self.mqtt_client.disconnect()
                self.is_connected = False
                print("已断开MQTT连接")
            except Exception as e:
                print(f"断开连接时出错: {e}")

    def _on_message(self, topic, msg):
        """
        内部消息回调处理函数

        参数:
        topic -- 消息主题
        msg -- 消息内容
        """
        topic_str = topic.decode() if isinstance(topic, bytes) else topic
        msg_str = msg.decode() if isinstance(msg, bytes) else msg

        self.last_message = (topic_str, msg_str)

        # 调用针对该主题的回调函数
        if topic_str in self.subscriptions and self.subscriptions[topic_str]:
            try:
                self.subscriptions[topic_str](topic_str, msg_str)
            except Exception as e:
                print(f"回调函数执行错误: {e}")

    def publish(self, topic, message, retain=False, qos=0):
        """
        发布消息到指定主题

        参数:
        topic -- 主题名称
        message -- 消息内容，可以是字符串或字典(会被转换为JSON)
        retain -- 是否保留消息
        qos -- 服务质量等级(0, 1, 2)

        返回:
        bool -- 是否成功发布
        """
        if not self.is_connected:
            print("MQTT未连接，无法发布消息")
            return False

        try:
            # 将字典转换为JSON字符串
            if isinstance(message, dict):
                message = json.dumps(message)

            # 确保消息是字节类型
            if isinstance(message, str):
                message = message.encode()

            # 确保主题是字节类型
            if isinstance(topic, str):
                topic = topic.encode()

            self.mqtt_client.publish(topic, message, retain, qos)
            return True

        except Exception as e:
            print(f"发布消息失败: {e}")
            self.is_connected = False
            return False

    def subscribe(self, topic, callback=None):
        """
        订阅主题

        参数:
        topic -- 主题名称
        callback -- 回调函数，接收两个参数(topic, msg)

        返回:
        bool -- 是否成功订阅
        """
        if not self.is_connected:
            print("MQTT未连接，无法订阅主题")
            self.subscriptions[topic] = callback
            return False

        try:
            # 存储回调函数
            self.subscriptions[topic] = callback

            # 订阅主题
            self.mqtt_client.subscribe(topic.encode() if isinstance(topic, str) else topic)
            print(f"已订阅主题: {topic}")
            return True

        except Exception as e:
            print(f"订阅主题失败: {e}")
            self.is_connected = False
            return False

    def check_msg(self):
        """
        检查是否有新消息
        需要在主循环中定期调用此方法
        """
        if self.is_connected:
            try:
                self.mqtt_client.check_msg()
            except Exception as e:
                print(f"检查消息时出错: {e}")
                self.is_connected = False

    def wait_msg(self):
        """
        阻塞等待消息
        注意：此方法会阻塞程序执行，直到收到消息
        """
        if self.is_connected:
            try:
                self.mqtt_client.wait_msg()
            except Exception as e:
                print(f"等待消息时出错: {e}")
                self.is_connected = False

    def reconnect(self):
        """重新连接到MQTT服务器"""
        self.disconnect()
        return self.connect()